if __name__ == '__main__':
    with open('tasks_data.txt', 'r', encoding='utf-8') as f:
        datas = f.readlines()
    train_data = make_data(datas)
    train_num_data = [[word2id[word] for word in line]   for line in train_data]
    batch_size = 16  # 64
    epochs = 8
    dataset = MyDataSet(train_num_data)
    data_loader = Data.DataLoader(dataset,batch_size=batch_size, collate_fn=dataset.padding_batch)
    model = Transformer().to(device)
    model.load_state_dict(torch.load('chat_model.pt')) # 加载预训练模型
    train(model, data_loader) # 微调训练
    torch.save(model.state_dict(),'fine_tuning_model.pt')  # 保存微调模型 


 